import torch
import numpy as np
from auxmodels.utils.calc import np_sigmoid, np_max, xywh2xyxy

from ..utils.cal_funcs import decode_bbox


class Decoder:
    """模型解码"""

    def __init__(self, thres_conf=0.5, stride=4, scale=64) -> None:
        self.thres_conf = thres_conf
        self.stride = stride
        self.scale = scale

    def run(self, preds):
        if isinstance(preds, tuple):
            preds_type = preds[0]
        else:
            preds_type = preds

        if torch.is_tensor(preds_type):
            f_where = torch.where
            f_sigmoid = torch.sigmoid
            f_cat = torch.cat
            f_stack = torch.stack
            f_tanh = torch.tanh
            f_max = torch.max
        else:
            f_where = np.where
            f_sigmoid = np_sigmoid
            f_cat = np.concatenate
            f_stack = np.stack
            f_tanh = np.tanh
            f_max = np_max

        out_bbox, out_hm, out_pool = preds
        batch_size = out_hm.shape[0]
        crest = out_hm == out_pool

        b, c, row, col = f_where(crest)
        hm_activated = f_sigmoid(out_hm[b, c, row, col])
        obj_mask = hm_activated > self.thres_conf
        b, c, row, col, conf = (
            b[obj_mask],
            c[obj_mask],
            row[obj_mask],
            col[obj_mask],
            hm_activated[obj_mask],
        )

        if len(conf) == 0:
            return [None for _ in range(batch_size)]

        out_box = decode_bbox(
            f_tanh(out_bbox[b, :, row, col]),
            row,
            col,
            self.stride,
            self.stride * self.scale,
            f_stack,
        )

        obj_bbox = xywh2xyxy(out_box)
        obj_conf = f_max(f_sigmoid(out_hm[b, :, row, col]), -1)

        results = f_cat((f_stack((b, obj_conf[0], obj_conf[1]), -1), obj_bbox), -1)
        results_div = [results[results[:, 0] == ib] for ib in range(batch_size)]
        results_div = [x if len(x) > 0 else None for x in results_div]
        return results_div
